# models/efficient_time_llm.py
import paddle
import paddle.nn as nn

from models.time_llm_block import TimeLLMBlock


class EfficientTimeLLM(nn.Layer):
    def __init__(self, seq_len=96, pred_len=24, d_model=128, num_layers=3, top_k=5, n_vars=1):
        super().__init__()
        self.seq_len = seq_len
        self.pred_len = pred_len
        self.n_vars = n_vars  # ✅ 新增：支持多变量

        # 1. 多变量嵌入
        self.value_embed = nn.Linear(n_vars, d_model)  # 输入从 1 → n_vars

        # 2. 时间特征嵌入
        self.hour_embed = nn.Embedding(24, d_model)
        self.dow_embed = nn.Embedding(7, d_model)

        # 3. 位置嵌入
        self.position_embed = nn.Embedding(seq_len, d_model)

        # 4. 编码器
        self.encoder = nn.LayerList([
            TimeLLMBlock(d_model, top_k) for _ in range(num_layers)
        ])

        # 5. 解码器
        self.decoder = nn.LayerList([
            TimeLLMBlock(d_model, top_k) for _ in range(2)
        ])

        # 6. 预测头
        self.head = nn.Sequential(
            nn.Linear(d_model, d_model * 2),
            nn.GELU(),
            nn.Linear(d_model * 2, pred_len)
        )

    def forward(self, X, time_features):
        """
        X: [B, L, N] - 多变量输入
        time_features: [B, L, 2] - 时间特征（hour, dow）
        """
        if not isinstance(X, paddle.Tensor):
            X = paddle.to_tensor(X, dtype='float32')

        # 数值嵌入
        value_emb = self.value_embed(X)  # [B, L, d_model]

        # 时间特征嵌入
        hour_emb = self.hour_embed(time_features[:, :, 0].astype('int64'))
        dow_emb = self.dow_embed(time_features[:, :, 1].astype('int64'))
        time_emb = hour_emb + dow_emb

        # 位置嵌入
        positions = paddle.arange(self.seq_len).unsqueeze(0).tile([X.shape[0], 1])
        pos_emb = self.position_embed(positions)

        # 联合嵌入
        x = value_emb + time_emb + pos_emb

        # 编码器
        seasonal_list = []
        trend_accum = None

        for layer in self.encoder:
            trend, seasonal = layer(x)
            seasonal_list.append(seasonal)
            trend_accum = trend if trend_accum is None else trend_accum + trend

        seasonal_avg = sum(seasonal_list) / len(seasonal_list)

        # 解码器
        for layer in self.decoder:
            _, seasonal_avg = layer(seasonal_avg)

        # 预测
        return self.head(seasonal_avg[:, -self.pred_len:])
